from __future__ import absolute_import, division, print_function

import tensorflow as tf


def parse_record(record):
    features = {
        'survived': tf.io.FixedLenFeature([], tf.int64),
        'pclass': tf.io.FixedLenFeature([], tf.int64),
        'name': tf.io.FixedLenFeature([], tf.string),
        'sex': tf.io.FixedLenFeature([], tf.string),
        'age': tf.io.FixedLenFeature([], tf.float32),
        'sibsp': tf.io.FixedLenFeature([], tf.int64),
        'parch': tf.io.FixedLenFeature([], tf.int64),
        'ticket': tf.io.FixedLenFeature([], tf.string),
        'fare': tf.io.FixedLenFeature([], tf.float32),
    }
    return tf.io.parse_single_example(record, features=features)


if __name__ == '__main__':
    filenames = ['../data/titanic_dataset.tfrecord']
    data = tf.data.TFRecordDataset(filenames)
    map_data = data.map(parse_record)
    # Prefetch batch (pre-load batch for faster consumption).
    prefetch_data = map_data.prefetch(buffer_size=1)
    iterator = prefetch_data.make_initializable_iterator()
    with tf.Session() as sess:
        sess.run(iterator.initializer)
        x = iterator.get_next()
        for i in range(3):
            print(sess.run(x))
